import numpy as np
import pandas as pd
import os
import pickle
import h5py
import random
import matplotlib.pyplot as plt
import sys
from config import *
from utils import *
from model import Logistic, DoubleLogistic, Linear, SimpleLogistic, NN, NaiveBayes, NoduleLogistic


WDModelDir = '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch1/results/'
JYModelDir = '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch2/result/'


testsetNameList1 = [
				# WDModelDir + 'vgg_shortcut_cddv15_v2_without_test_epoch-0089_stage2/vgg_shortcut_cddv15_v2_without_test_merge',
				WDModelDir + 'vgg_shortcut_cddv15_v2_without_test_epoch-0099_stage2/vgg_shortcut_cddv15_v2_without_test_merge',
				# WDModelDir + 'vgg_shortcut_cddv9_without_test_epoch-0089_stage2/vgg_shortcut_cddv9_without_test_merge',
				# WDModelDir + 'vgg_shortcut_cddv9_without_test_epoch-0099_stage2/vgg_shortcut_cddv9_without_test_merge',
				
				# JYModelDir + 'v15_without_test_90/vgg13_v15_merge',
				# JYModelDir + 'v15_without_test_97/vgg13_v15_merge',
				# JYModelDir + 'v15_without_test_99/vgg13_v15_merge',
				# JYModelDir + 'v9_without_test_120/vgg13_v9_merge',
				]

testsetNameList2 = [
				'/ssd_1t/wangd/kaggle_results/vgg_shortcut_cddv15_v2_without_test_testset_merge',
				# '/ssd_1t/wangd/kaggle_results/vgg_shortcut_cddv9_without_test_testset_merge',
				# '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch1/results/vgg_shortcut_cddv9_without_test/vgg_shortcut_cddv9_without_test_merge',
				# '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch2/result(baseline)/v15_without_test_90/vgg13_v15_merge',
				# '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch2/result(baseline)/v15_without_test_97/vgg13_v15_merge',
				# '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch2/result(baseline)/v15_without_test_99/vgg13_v15_merge',
				# '/data_4t/Kaggle/submission_code_data/data/tmp_classification_branch2/result(baseline)/v9_without_test_120/vgg13_v9_merge',
				]

submissionDir = '/data_4t/Kaggle/submission_code_data/data/submission/'

def testAll1(trainsetData):
	#########################################################################################
	# Test on the Kaggle testSet
	# load test data
	testsetData = []
	for datasetName in testsetNameList1:
		print(datasetName)
		dataset = dict()
		dataset['name'] = datasetName
		dataset['predScan'] = []
		for iFold in range(nFold):
			fName = datasetName + str(iFold) + '.pkl'
			# print(fName)
			f = open(fName, 'rb')
			predCdds = pickle.load(f)
			# A list of predict scan at different fold
			dataset['predScan'].append(cdds2scan(predCdds, mode='stage2'))
		testsetData.append(dataset)

	testsetScanList = list(testsetData[0]['predScan'][0].keys())
	predScanScore = dict()
	for scanid in testsetScanList: predScanScore[scanid] = []

	##########################################################################################
	# Predict
	for i, testset in enumerate(testsetData):
		if averageFoldFirst:
			predScan = dict()
			for scanid in testsetScanList:
				avgScan = AvgOverScan(testset['predScan'], scanid)
				predScan[scanid] = avgScan
			trainsetData[i]['allmodel'].dataSet = predScan
			for scanid in testsetScanList:
				score = trainsetData[i]['allmodel'].predictAtName(scanid)
				predScanScore[scanid].append(score)
		else:
			predScan = dict()
			for iFold in range(nFold):
				for scanid in testsetScanList:
					predScan[scanid] = testset['predScan'][iFold][scanid]
				trainsetData[i]['allmodel'].dataSet = predScan
				for scanid in testsetScanList:
					# print(scanid)
					score = trainsetData[i]['allmodel'].predictAtName(scanid)
					predScanScore[scanid].append(score)


	###########################################################################################
	# Output submission data
	submissionData = dict()
	for (scanid, predScore) in predScanScore.items():
		s = np.sum(predScore)
		s = s - np.max(predScore) - np.min(predScore)
		score = s / (len(predScore) - 2)
		# score = np.mean(predScore)
		submissionData[scanid] = score
		

	output_f = open(submissionDir+'stage2_submission1.csv', 'w')
	scanidList = sorted(submissionData.keys())
	output_f.write('id,cancer\n')
	for scanid in scanidList:
		output_f.write('{},{}\n'.format(scanid, submissionData[scanid]))
	output_f.close()

def testAll2(trainsetData):
	#########################################################################################
	# Test on the Kaggle testSet
	# load test data
	testsetData = []
	for datasetName in testsetNameList2:
		print(datasetName)
		dataset = dict()
		dataset['name'] = datasetName
		dataset['predScan'] = []
		for iFold in range(nFold):
			fName = datasetName + str(iFold) + '.pkl'
			# print(fName)
			f = open(fName, 'rb')
			predCdds = pickle.load(f)
			# A list of predict scan at different fold
			dataset['predScan'].append(cdds2scan(predCdds, mode='stage2'))
		testsetData.append(dataset)

	testsetScanList = list(testsetData[0]['predScan'][0].keys())
	predScanScore = dict()
	for scanid in testsetScanList: predScanScore[scanid] = []

	##########################################################################################
	# Predict
	for i, testset in enumerate(testsetData):
		if averageFoldFirst:
			predScan = dict()
			for scanid in testsetScanList:
				avgScan = AvgOverScan(testset['predScan'], scanid)
				predScan[scanid] = avgScan
			trainsetData[i]['allmodel'].dataSet = predScan
			for scanid in testsetScanList:
				score = trainsetData[i]['allmodel'].predictAtName(scanid)
				predScanScore[scanid].append(score)
		else:
			predScan = dict()
			for iFold in range(nFold):
				for scanid in testsetScanList:
					predScan[scanid] = testset['predScan'][iFold][scanid]
				trainsetData[i]['allmodel'].dataSet = predScan
				for scanid in testsetScanList:
					# print(scanid)
					score = trainsetData[i]['allmodel'].predictAtName(scanid)
					predScanScore[scanid].append(score)


	###########################################################################################
	# Output submission data
	submissionData = dict()
	for (scanid, predScore) in predScanScore.items():
		s = np.sum(predScore)
		s = s - np.max(predScore) - np.min(predScore)
		score = s / (len(predScore) - 2)
		# score = np.mean(predScore)
		submissionData[scanid] = score
		

	output_f = open(submissionDir+'stage2_submission2.csv', 'w')
	scanidList = sorted(submissionData.keys())
	output_f.write('id,cancer\n')
	for scanid in scanidList:
		output_f.write('{},{}\n'.format(scanid, submissionData[scanid]))
	output_f.close()

if __name__ == '__main__':
	print(modelType)

	f = open('trainsetData1.pkl', 'rb')
	trainsetData = pickle.load(f)
	f.close()
	testAll1(trainsetData[iset:])

	f = open('trainsetData2.pkl', 'rb')
	trainsetData = pickle.load(f)
	f.close()
	testAll2(trainsetData[iset:])
